'''
* This is the projet for brtc R&D  Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description MCP服务器实现模块，提供上下文管理的服务端功能
* @File: mcp_server.py
* @Time: 2025/01/27 14:30:00
* @All Rights Reserve By Brtc
'''

import json
import socket
import threading
from datetime import datetime
from typing import Dict, List, Optional, Any
from copy import deepcopy

from .mcp_protocol import (
    MCPMessage,
    MCPResponse,
    Context,
    OperationType,
    StatusCode,
    Tool,
    ToolParameter
)


class MCPServer:
    """MCP服务器类"""
    
    def __init__(self, host: str = "localhost", port: int = 8888):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化MCP服务器
        * @Date :2025/01/27 14:30:00
        * @Param: host:服务器主机地址, port:服务器端口号
        * @return: 无
        """
        self.host = host
        self.port = port
        self.socket = None
        self.contexts: Dict[str, Context] = {}
        self.tools: Dict[str, Tool] = {}
        self.running = False
        self.lock = threading.Lock()
    
    def start(self):
        """
        * @Author Leon-liao
        * @Function: start
        * @Description //启动MCP服务器
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 无
        """
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.socket.bind((self.host, self.port))
        self.socket.listen(5)
        self.running = True
        
        print(f"MCP服务器已启动，监听 {self.host}:{self.port}")
        
        while self.running:
            try:
                client_socket, address = self.socket.accept()
                print(f"新客户端连接: {address}")
                client_thread = threading.Thread(
                    target=self._handle_client,
                    args=(client_socket, address)
                )
                client_thread.daemon = True
                client_thread.start()
            except Exception as e:
                if self.running:
                    print(f"接受连接时出错: {e}")
    
    def stop(self):
        """
        * @Author Leon-liao
        * @Function: stop
        * @Description //停止MCP服务器
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 无
        """
        self.running = False
        if self.socket:
            self.socket.close()
        print("MCP服务器已停止")
    
    def _handle_client(self, client_socket: socket.socket, address: tuple):
        """
        * @Author Leon-liao
        * @Function: _handle_client
        * @Description //处理客户端请求
        * @Date :2025/01/27 14:30:00
        * @Param: client_socket:客户端套接字, address:客户端地址
        * @return: 无
        """
        try:
            while self.running:
                data = client_socket.recv(4096)
                if not data:
                    break
                
                try:
                    message_dict = json.loads(data.decode('utf-8'))
                    message = MCPMessage.from_dict(message_dict)
                    response = self._process_message(message)
                    
                    response_data = json.dumps(response.to_dict()).encode('utf-8')
                    client_socket.sendall(response_data)
                except json.JSONDecodeError:
                    response = MCPResponse(
                        StatusCode.BAD_REQUEST,
                        "无效的JSON格式",
                        request_id=message_dict.get("header", {}).get("message_id")
                    )
                    client_socket.sendall(
                        json.dumps(response.to_dict()).encode('utf-8')
                    )
                except Exception as e:
                    response = MCPResponse(
                        StatusCode.INTERNAL_ERROR,
                        f"处理请求时出错: {str(e)}",
                        request_id=message_dict.get("header", {}).get("message_id")
                    )
                    client_socket.sendall(
                        json.dumps(response.to_dict()).encode('utf-8')
                    )
        except Exception as e:
            print(f"处理客户端 {address} 时出错: {e}")
        finally:
            client_socket.close()
            print(f"客户端 {address} 已断开连接")
    
    def _process_message(self, message: MCPMessage) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _process_message
        * @Description //处理MCP消息并返回响应
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象
        * @return: 返回MCP响应对象
        """
        operation = OperationType(message.context["operation"])
        context_id = message.context["id"]
        request_id = message.header["message_id"]
        
        try:
            if operation == OperationType.CREATE:
                return self._handle_create(message, request_id)
            elif operation == OperationType.READ:
                return self._handle_read(context_id, request_id)
            elif operation == OperationType.UPDATE:
                return self._handle_update(message, request_id)
            elif operation == OperationType.DELETE:
                return self._handle_delete(context_id, request_id)
            elif operation == OperationType.SEARCH:
                return self._handle_search(message, request_id)
            elif operation == OperationType.COPY:
                return self._handle_copy(message, request_id)
            elif operation == OperationType.LIST_TOOLS:
                return self._handle_list_tools(request_id)
            elif operation == OperationType.CALL_TOOL:
                return self._handle_call_tool(message, request_id)
            else:
                return MCPResponse(
                    StatusCode.BAD_REQUEST,
                    f"不支持的操作类型: {operation.value}",
                    request_id=request_id
                )
        except Exception as e:
            return MCPResponse(
                StatusCode.INTERNAL_ERROR,
                f"处理操作时出错: {str(e)}",
                request_id=request_id
            )
    
    def _handle_create(self, message: MCPMessage, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_create
        * @Description //处理创建上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            context_id = message.context["id"]
            if context_id in self.contexts:
                return MCPResponse(
                    StatusCode.CONFLICT,
                    f"上下文 {context_id} 已存在",
                    request_id=request_id
                )
            
            context = Context(
                context_id=context_id,
                content=message.payload.get("content"),
                metadata=message.payload.get("metadata", {}),
                extensions=message.payload.get("extensions", {})
            )
            
            self.contexts[context_id] = context
            
            return MCPResponse(
                StatusCode.CREATED,
                "上下文创建成功",
                data={"context": context.to_dict()},
                request_id=request_id
            )
    
    def _handle_read(self, context_id: str, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_read
        * @Description //处理读取上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            if context_id not in self.contexts:
                return MCPResponse(
                    StatusCode.NOT_FOUND,
                    f"上下文 {context_id} 不存在",
                    request_id=request_id
                )
            
            context = self.contexts[context_id]
            return MCPResponse(
                StatusCode.SUCCESS,
                "读取上下文成功",
                data={"context": context.to_dict()},
                request_id=request_id
            )
    
    def _handle_update(self, message: MCPMessage, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_update
        * @Description //处理更新上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            context_id = message.context["id"]
            if context_id not in self.contexts:
                return MCPResponse(
                    StatusCode.NOT_FOUND,
                    f"上下文 {context_id} 不存在",
                    request_id=request_id
                )
            
            context = self.contexts[context_id]
            
            # 更新内容
            if "content" in message.payload:
                context.content = message.payload["content"]
            
            # 更新元数据
            if "metadata" in message.payload:
                context.metadata.update(message.payload["metadata"])
            
            # 更新扩展字段
            if "extensions" in message.payload:
                context.extensions.update(message.payload["extensions"])
            
            context.update_modified_time()
            
            return MCPResponse(
                StatusCode.SUCCESS,
                "上下文更新成功",
                data={"context": context.to_dict()},
                request_id=request_id
            )
    
    def _handle_delete(self, context_id: str, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_delete
        * @Description //处理删除上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            if context_id not in self.contexts:
                return MCPResponse(
                    StatusCode.NOT_FOUND,
                    f"上下文 {context_id} 不存在",
                    request_id=request_id
                )
            
            del self.contexts[context_id]
            
            return MCPResponse(
                StatusCode.SUCCESS,
                "上下文删除成功",
                request_id=request_id
            )
    
    def _handle_search(self, message: MCPMessage, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_search
        * @Description //处理搜索上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            search_criteria = message.payload.get("criteria", {})
            results = []
            
            for context_id, context in self.contexts.items():
                match = True
                
                # 按标签搜索
                if "labels" in search_criteria:
                    context_labels = context.metadata.get("labels", [])
                    required_labels = search_criteria["labels"]
                    if not all(label in context_labels for label in required_labels):
                        match = False
                
                # 按作用域搜索
                if "scope" in search_criteria:
                    if context.metadata.get("scope") != search_criteria["scope"]:
                        match = False
                
                # 按内容关键词搜索
                if "keyword" in search_criteria:
                    keyword = search_criteria["keyword"].lower()
                    content_str = str(context.content).lower()
                    if keyword not in content_str:
                        match = False
                
                if match:
                    results.append(context.to_dict())
            
            return MCPResponse(
                StatusCode.SUCCESS,
                f"搜索完成，找到 {len(results)} 个匹配的上下文",
                data={"results": results, "count": len(results)},
                request_id=request_id
            )
    
    def _handle_copy(self, message: MCPMessage, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_copy
        * @Description //处理复制上下文请求
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            source_id = message.context["id"]
            target_id = message.payload.get("target_id")
            
            if not target_id:
                return MCPResponse(
                    StatusCode.BAD_REQUEST,
                    "缺少目标上下文ID",
                    request_id=request_id
                )
            
            if source_id not in self.contexts:
                return MCPResponse(
                    StatusCode.NOT_FOUND,
                    f"源上下文 {source_id} 不存在",
                    request_id=request_id
                )
            
            if target_id in self.contexts:
                return MCPResponse(
                    StatusCode.CONFLICT,
                    f"目标上下文 {target_id} 已存在",
                    request_id=request_id
                )
            
            # 深拷贝源上下文
            source_context = self.contexts[source_id]
            new_context = Context(
                context_id=target_id,
                content=deepcopy(source_context.content),
                metadata=deepcopy(source_context.metadata),
                extensions=deepcopy(source_context.extensions)
            )
            
            # 更新创建时间
            new_context.metadata["created"] = datetime.utcnow().isoformat() + "Z"
            new_context.metadata["modified"] = datetime.utcnow().isoformat() + "Z"
            
            self.contexts[target_id] = new_context
            
            return MCPResponse(
                StatusCode.CREATED,
                "上下文复制成功",
                data={"context": new_context.to_dict()},
                request_id=request_id
            )
    
    def register_tool(self, tool: Tool):
        """
        * @Author Leon-liao
        * @Function: register_tool
        * @Description //注册工具到服务器
        * @Date :2025/01/27 14:30:00
        * @Param: tool:工具对象
        * @return: 无
        """
        with self.lock:
            self.tools[tool.name] = tool
            print(f"工具 '{tool.name}' 已注册")
    
    def _handle_list_tools(self, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_list_tools
        * @Description //处理列出所有工具的请求
        * @Date :2025/01/27 14:30:00
        * @Param: request_id:请求ID
        * @return: 返回MCP响应对象
        """
        with self.lock:
            tools_list = [tool.to_dict() for tool in self.tools.values()]
            return MCPResponse(
                StatusCode.SUCCESS,
                f"找到 {len(tools_list)} 个可用工具",
                data={"tools": tools_list},
                request_id=request_id
            )
    
    def _handle_call_tool(self, message: MCPMessage, request_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _handle_call_tool
        * @Description //处理调用工具的请求
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象, request_id:请求ID
        * @return: 返回MCP响应对象
        """
        tool_name = message.payload.get("tool_name")
        if not tool_name:
            return MCPResponse(
                StatusCode.BAD_REQUEST,
                "缺少工具名称",
                request_id=request_id
            )
        
        with self.lock:
            if tool_name not in self.tools:
                return MCPResponse(
                    StatusCode.NOT_FOUND,
                    f"工具 '{tool_name}' 不存在",
                    request_id=request_id
                )
            
            tool = self.tools[tool_name]
            
            # 检查工具是否有处理函数
            if not tool.handler:
                return MCPResponse(
                    StatusCode.INTERNAL_ERROR,
                    f"工具 '{tool_name}' 没有注册处理函数",
                    request_id=request_id
                )
            
            # 验证参数
            tool_args = message.payload.get("arguments", {})
            missing_params = []
            for param in tool.parameters:
                if param.required and param.name not in tool_args:
                    if param.default is None:
                        missing_params.append(param.name)
            
            if missing_params:
                return MCPResponse(
                    StatusCode.BAD_REQUEST,
                    f"缺少必需参数: {', '.join(missing_params)}",
                    request_id=request_id
                )
            
            # 准备参数（包括默认值）
            final_args = {}
            for param in tool.parameters:
                if param.name in tool_args:
                    final_args[param.name] = tool_args[param.name]
                elif param.default is not None:
                    final_args[param.name] = param.default
            
            # 调用工具处理函数
            try:
                result = tool.handler(**final_args)
                return MCPResponse(
                    StatusCode.SUCCESS,
                    f"工具 '{tool_name}' 执行成功",
                    data={"result": result},
                    request_id=request_id
                )
            except Exception as e:
                return MCPResponse(
                    StatusCode.INTERNAL_ERROR,
                    f"工具执行出错: {str(e)}",
                    request_id=request_id
                )
